"""
base_llm.py 模块
大语言模型的抽象基类定义（优化后）
"""
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Union


class BaseLLM(ABC):
    """大语言模型基类，定义核心接口（简化抽象方法要求）"""

    def __init__(
        self,
        model_name: str,
        temperature: float = 0.7,
        max_tokens: int = 1024,
        **kwargs
    ):
        """
        初始化LLM（非抽象方法，子类可直接继承或重写）

        Args:
            model_name: 模型名称
            temperature: 温度参数，控制生成文本的随机性（0-1）
            max_tokens: 生成文本的最大长度
            **kwargs: 其他模型参数（如设备、量化配置等）
        """
        self.model_name = model_name
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.kwargs = kwargs  # 存储额外参数


    @abstractmethod
    def generate(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        """
        核心生成接口（抽象方法，子类必须实现）
        单轮文本生成，基于输入提示词返回结果

        Args:
            prompt: 输入提示词（字符串）
            stop: 停止符列表（遇到这些字符则停止生成）

        Returns:
            生成的文本字符串
        """
        pass


    def chat(
        self,
        messages: List[Dict[str, str]],
        stop: Optional[List[str]] = None
    ) -> str:
        """
        多轮对话接口（非抽象方法，提供默认实现）
        支持多轮消息历史，默认将消息转换为prompt调用generate

        Args:
            messages: 消息列表，格式为 [{"role": "user/assistant", "content": "文本"}]
            stop: 停止符列表

        Returns:
            生成的回复文本
        """
        # 默认将多轮消息转换为单轮prompt（适用于不支持原生多轮的模型）
        prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
        prompt += "\nassistant: "  # 引导模型生成回复
        return self.generate(prompt, stop=stop)